"""
============================================================
📦 File: checkpoint_manager.py
🧠 Description:
    本模块实现模型训练过程中的 Checkpoint 管理器，用于：
        - 保存模型、优化器、学习率调度器等状态；
        - 加载已保存的 checkpoint；
        - 自动保存最优模型。

✨ Features:
    ✅ 定期保存模型 (按 epoch)
    ✅ 自动保存最优模型 (基于验证集指标)
    ✅ 支持从 checkpoint 恢复训练状态
    ✅ 支持与自定义配置类 CheckpointConfig 兼容

📌 Usage:
    >>> manager = CheckpointManager(model, optimizer, scheduler, config.checkpoint)
    >>> manager.save_checkpoint(epoch, val_metric, is_best=True)
    >>> start_epoch = manager.load_checkpoint(config.resume_path)

🚀 Author: kkli
📅 Date: 2025-10-30
============================================================
"""

import os
from typing import Optional

import torch
import logging

from config_manager import CheckpointConfig


class CheckpointManager:
    """
    管理模型的保存与加载。

    Args:
        trainer_manager (TrainerManager): 训练管理器
        config (CheckpointConfig): 模型保存配置类
    """

    def __init__(self, trainer_manager, config:CheckpointConfig):
        """初始化 Checkpoint 管理器"""
        self.trainer_manager = trainer_manager
        self.model = trainer_manager.model
        self.optimizer = trainer_manager.optimizer
        self.scheduler = trainer_manager.scheduler
        self.config = config

        # 模型保存路径，例如：./checkpoints/resnet50/
        self.save_dir = os.path.join(self.config.save_dir, getattr(self.trainer_manager.model_manager.config, "model_name", "model"))
        os.makedirs(self.save_dir, exist_ok=True)

        # 用于追踪最优指标
        self.best_metric_value = None
        self.is_bast = False

    def should_save(self, epoch: int, val_metric: Optional[dict] = None) -> bool:
        """
        判断当前是否需要保存 checkpoint。
        """
        # 如果只保存最优模型
        if self.config.save_best_only:
            if val_metric is not None:
                main_metric = val_metric.get(self.config.main_metric, None)
                if main_metric is not None:
                    if self.best_metric_value is None or main_metric > self.best_metric_value:
                        self.best_metric_value = main_metric
                        logging.info(f"🏆 新的最优模型（{self.config.main_metric}={main_metric:.4f}）")
                        self.is_bast=True
                        return True
                else:
                    logging.warning("⚠️ 未提供验证指标 val_metric，无法判断最优模型，跳过保存。")
            else:
                logging.warning("⚠️ 未提供验证指标 val_metric，无法判断最优模型，跳过保存。")
            return False

        # 否则，按 save_interval 保存
        return (epoch + 1) % self.config.save_interval == 0

        # -----------------------------------------------------

    def save_checkpoint(self, epoch: int, val_metric: Optional[dict] = None):
        """
        保存模型训练状态。

        Args:
            epoch (int): 当前 epoch
            val_metric (dict, optional): 当前验证指标
        """
        # 判断是否需要保存
        if not self.should_save(epoch, val_metric):
            return

        # 构建保存字典
        state = {
            "epoch": epoch,
            "model_state": self.model.state_dict(),
            "optimizer_state": self.optimizer.state_dict(),
            "scheduler_state": self.scheduler.state_dict() if self.scheduler else None,
            "val_metric": val_metric,
        }

        if self.is_bast:
            pt_path = os.path.join(self.save_dir, "best_model.pth")
            torch.save(state, pt_path)
            logging.info(f"💾 最优模型已保存至 {pt_path}")
            self.is_bast=False

        if not self.config.save_best_only:
            pt_path = os.path.join(self.save_dir, f"epoch_{epoch + 1}.pth")
            torch.save(state, pt_path)
            logging.info(f"💾 模型已保存至 {pt_path}")


    def load_checkpoint(self, path: Optional[str] = None) -> int:
        """
        加载模型 checkpoint。

        Args:
            path (str, optional): checkpoint 文件路径；
                若未提供则尝试使用 config.resume_path。
        Returns:
            int: checkpoint 中保存的 epoch 编号
        """
        path = path or getattr(self.config, "resume_path", None)
        if path is None or not os.path.exists(path):
            logging.warning("⚠️ 未提供有效的 checkpoint 路径，跳过加载。")
            return 0

        ckpt = torch.load(path, map_location=getattr(self.config, "device", "cpu"))

        # 恢复模型与优化器状态
        self.model.load_state_dict(ckpt["model_state"])
        if "optimizer_state" in ckpt and self.optimizer:
            self.optimizer.load_state_dict(ckpt["optimizer_state"])
        if "scheduler_state" in ckpt and self.scheduler:
            self.scheduler.load_state_dict(ckpt["scheduler_state"])

        self.best_metric_value = ckpt.get("val_metric", None)
        start_epoch = ckpt.get("epoch", 0)

        logging.info(f"✅ 已从 {path} 恢复模型，起始 epoch: {start_epoch}")
        return start_epoch
